#!/usr/bin/env python3
"""
策略模型相关函数
用于H1机器人的平衡测试
"""
import os
import torch
import numpy as np
import math

def init_policy(policy_config):
    """
    初始化策略模型
    
    Args:
        policy_config: 策略模型配置
        
    Returns:
        policy: 加载的策略模型，如果加载失败则返回None
    """
    print("准备加载策略模型...")
    try:
        # 尝试加载策略模型
        policy_path = policy_config["policy_path"]
        if os.path.exists(policy_path):
            policy = torch.jit.load(policy_path, map_location=policy_config["device"])
            policy.to(policy_config["device"])
            print(f"策略模型成功加载：{policy_path}")
            return policy
        else:
            print(f"找不到策略模型文件：{policy_path}")
            return None
    except Exception as e:
        print(f"加载策略模型时发生错误：{e}")
        return None

def create_observation_buffer(policy_config):
    """
    创建观察历史缓冲区
    
    Args:
        policy_config: 策略模型配置
        
    Returns:
        obs_history_buf: 观察历史缓冲区
    """
    # 创建观察历史缓冲区
    obs_context_len = policy_config["obs_context_len"]
    expected_obs_size = policy_config["expected_obs_size"]
    device = policy_config["device"]
    
    # 初始化为零矩阵
    obs_history_buf = torch.zeros((1, obs_context_len, expected_obs_size), device=device)
    return obs_history_buf

def get_observation(model, data, home_qpos, policy_config):
    """
    根据当前状态构建观察向量
    
    Args:
        model: MuJoCo模型对象
        data: MuJoCo数据对象
        home_qpos: 默认关节位置
        policy_config: 策略模型配置
        
    Returns:
        obs_buf: 观察向量
    """
    # 创建观察向量
    expected_obs_size = policy_config["expected_obs_size"]
    device = policy_config["device"]
    
    # 初始化观察数据数组
    obs_data = np.zeros(expected_obs_size)
    
    # 1. 基本姿态信息（方向和角速度）
    # 获取基座四元数
    base_quat = data.qpos[3:7].copy()
    
    # 计算正确的方向（roll和pitch）
    # 简化版，假设四元数是标准化的
    roll = math.atan2(2.0 * (base_quat[0] * base_quat[1] + base_quat[2] * base_quat[3]),
                     1.0 - 2.0 * (base_quat[1] * base_quat[1] + base_quat[2] * base_quat[2]))
    pitch = math.asin(2.0 * (base_quat[0] * base_quat[2] - base_quat[3] * base_quat[1]))
    
    # 获取角速度
    ang_vel = data.qvel[3:6].copy()  # 角速度在浮动基座的速度中是3-5索引
    
    # 将姿态信息放入观察向量
    obs_data[0] = roll
    obs_data[1] = pitch
    obs_data[2:5] = ang_vel
    
    # 2. 关节位置和速度
    # 假设我们有19个关节（跳过浮动基座的7个自由度）
    num_joints = min(model.nq - 7, 19)  # 最多19个关节
    
    # 获取关节位置和速度
    joint_pos = data.qpos[7:7+num_joints].copy()
    joint_vel = data.qvel[6:6+num_joints].copy()  # 速度索引从6开始（浮动基座有6个速度自由度）
    
    # 计算关节位置相对于默认位置的偏移
    default_joint_pos = home_qpos[7:7+num_joints]
    joint_pos_offset = joint_pos - default_joint_pos
    
    # 将关节信息放入观察向量
    # 假设观察向量中关节位置从索引5开始，关节速度紧随其后
    obs_data[5:5+num_joints] = joint_pos_offset
    obs_data[5+num_joints:5+2*num_joints] = joint_vel
    
    # 3. 命令信号（最后部分）
    # 添加命令信号（如期望的速度）- 这里我们设置一个简单的前进命令
    command_idx = 5 + 2*num_joints
    obs_data[command_idx:command_idx+3] = np.array([0.5, 0.0, 0.0])  # 例如：以0.5 m/s的速度前进
    
    # 转换为PyTorch张量
    obs_buf = torch.tensor(obs_data, dtype=torch.float, device=device).unsqueeze(0)
    
    return obs_buf

def update_observation_history(obs_history_buf, obs_buf):
    """
    更新观察历史缓冲区
    
    Args:
        obs_history_buf: 观察历史缓冲区
        obs_buf: 当前观察向量
        
    Returns:
        updated_obs_history_buf: 更新后的观察历史缓冲区
    """
    # 移除最早的观察，添加新的观察
    updated_obs_history_buf = torch.cat([
        obs_history_buf[:, 1:],
        obs_buf.unsqueeze(1)
    ], dim=1)
    
    return updated_obs_history_buf

def get_policy_actions(policy, obs_history_buf):
    """
    使用策略模型获取动作
    
    Args:
        policy: 策略模型
        obs_history_buf: 观察历史缓冲区
        
    Returns:
        actions: 策略模型输出的动作
    """
    with torch.no_grad():
        # 运行策略模型获取动作
        actions = policy(obs_history_buf.detach())
    return actions

def get_retarget_positions(policy_actions, home_qpos, policy_config):
    """
    根据策略动作获取目标关节位置
    
    Args:
        policy_actions: 策略模型输出的动作，形状为[1, 19]
        home_qpos: 默认关节位置
        policy_config: 策略模型配置
        
    Returns:
        target_positions: 目标关节位置
    """
    # 将策略动作从张量转换为numpy数组
    actions_np = policy_actions.cpu().numpy().flatten()
    
    # 创建默认位置的副本进行修改
    # 只复制非浮动基座的关节位置（跳过前7个）
    default_dof_pos = home_qpos[7:].copy()
    target_positions = default_dof_pos.copy()
    
    # 假设我们的任务是"stand"，我们控制腿部和腰部的前11个关节
    if policy_config["task"] == "stand":
        # 腿部和腰部关节索引
        leg_waist_indices = list(range(11))  # 前11个关节
        
        for i, idx in enumerate(leg_waist_indices):
            if i < len(actions_np):
                # 按照动作缩放因子缩放动作
                scaled_action = actions_np[i] * policy_config["scale_action"]
                # 将动作应用到默认位置
                target_positions[idx] = target_positions[idx] + scaled_action
    
    return target_positions
